// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/sorting_writer.h>

namespace kumo::pollux::dwio::common {
    SortingWriter::SortingWriter(
        std::unique_ptr<Writer> writer,
        std::unique_ptr<exec::SortBuffer> sortBuffer,
        vector_size_t maxOutputRowsConfig,
        uint64_t maxOutputBytesConfig,
        uint64_t outputTimeSliceLimitMs)
        : outputWriter_(std::move(writer)),
          maxOutputRowsConfig_(maxOutputRowsConfig),
          maxOutputBytesConfig_(maxOutputBytesConfig),
          finishTimeSliceLimitMs_(outputTimeSliceLimitMs),
          sortPool_(sortBuffer->pool()),
          canReclaim_(sortBuffer->canSpill()),
          sortBuffer_(std::move(sortBuffer)) {
        POLLUX_CHECK_GT(maxOutputRowsConfig_, 0);
        POLLUX_CHECK_GT(maxOutputBytesConfig_, 0);
        if (sortPool_->parent()->reclaimer() != nullptr) {
            sortPool_->setReclaimer(MemoryReclaimer::create(this));
        }
        setState(State::kRunning);
    }

    SortingWriter::~SortingWriter() {
        sortPool_->release();
    }

    void SortingWriter::write(const VectorPtr &data) {
        checkRunning();
        sortBuffer_->addInput(data);
    }

    void SortingWriter::flush() {
        checkRunning();
        outputWriter_->flush();
    }

    bool SortingWriter::finish() {
        const uint64_t startTimeMs = getCurrentTimeMs();
        SCOPE_EXIT {
            const uint64_t flushTimeMs = getCurrentTimeMs() - startTimeMs;
            if (flushTimeMs != 0) {
                RECORD_HISTOGRAM_METRIC_VALUE(
                    kMetricHiveSortWriterFinishTimeMs, flushTimeMs);
            }
        };
        if (isRunning()) {
            sortBuffer_->noMoreInput();
            setState(State::kFinishing);
        }
        if (sortBuffer_ == nullptr) {
            return true;
        }

        const auto maxOutputBatchRows = outputBatchRows();
        RowVectorPtr output{nullptr};
        do {
            if (getCurrentTimeMs() - startTimeMs > finishTimeSliceLimitMs_) {
                return false;
            }
            output = sortBuffer_->getOutput(maxOutputBatchRows);
            if (output != nullptr) {
                outputWriter_->write(output);
            }
        } while (output != nullptr);

        sortBuffer_.reset();
        sortPool_->release();
        return true;
    }

    void SortingWriter::close() {
        POLLUX_CHECK(isFinishing());
        setState(State::kClosed);
        POLLUX_CHECK_NULL(sortBuffer_);
        outputWriter_->close();
    }

    void SortingWriter::abort() {
        setState(State::kAborted);

        sortBuffer_.reset();
        sortPool_->release();
        outputWriter_->abort();
    }

    bool SortingWriter::canReclaim() const {
        return canReclaim_;
    }

    uint64_t SortingWriter::reclaim(
        uint64_t targetBytes,
        memory::MemoryReclaimer::Stats &stats) {
        if (!canReclaim_) {
            return 0;
        }

        if (!isRunning() && !isFinishing()) {
            KLOG(WARNING) << "Can't reclaim from a not running hive sort writer pool: "
                 << sortPool_->name() << ", state: " << state()
                 << "used memory: " << succinctBytes(sortPool_->usedBytes())
                 << ", reserved memory: "
                 << succinctBytes(sortPool_->reservedBytes());
            ++stats.numNonReclaimableAttempts;
            return 0;
        }
        POLLUX_CHECK_NOT_NULL(sortBuffer_);

        return memory::MemoryReclaimer::run(
            [&]() {
                int64_t reclaimedBytes{0}; {
                    memory::ScopedReclaimedBytesRecorder recorder(
                        sortPool_, &reclaimedBytes);
                    sortBuffer_->spill();
                    sortPool_->release();
                }
                return reclaimedBytes;
            },
            stats);
    }

    vector_size_t SortingWriter::outputBatchRows() {
        vector_size_t estimatedMaxOutputRows =
                std::numeric_limits<vector_size_t>::max();
        if (sortBuffer_->estimateOutputRowSize().has_value() &&
            sortBuffer_->estimateOutputRowSize().value() != 0) {
            const uint64_t maxOutputRows =
                    maxOutputBytesConfig_ / sortBuffer_->estimateOutputRowSize().value();
            if (UNLIKELY(maxOutputRows > std::numeric_limits<vector_size_t>::max())) {
                return maxOutputRowsConfig_;
            }

            estimatedMaxOutputRows = maxOutputRows;
        }
        return std::min(estimatedMaxOutputRows, maxOutputRowsConfig_);
    }

    std::unique_ptr<memory::MemoryReclaimer> SortingWriter::MemoryReclaimer::create(
        SortingWriter *writer) {
        return std::unique_ptr<memory::MemoryReclaimer>(new MemoryReclaimer(writer));
    }

    bool SortingWriter::MemoryReclaimer::reclaimableBytes(
        const memory::MemoryPool &pool,
        uint64_t &reclaimableBytes) const {
        POLLUX_CHECK_EQ(pool.name(), writer_->sortPool_->name());

        reclaimableBytes = 0;
        if (!writer_->canReclaim()) {
            return false;
        }
        reclaimableBytes = pool.usedBytes();
        return true;
    }

    uint64_t SortingWriter::MemoryReclaimer::reclaim(
        memory::MemoryPool *pool,
        uint64_t targetBytes,
        uint64_t /*unused*/,
        memory::MemoryReclaimer::Stats &stats) {
        POLLUX_CHECK_EQ(pool->name(), writer_->sortPool_->name());

        return writer_->reclaim(targetBytes, stats);
    }
} // namespace kumo::pollux::dwio::common
